///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/all_pair_jast.h
///
/// \brief   header file for the all pair jastrow coefficient object
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef FORMIC_FQMC_ALL_PAIR_JAST_HEADER
#define FORMIC_FQMC_ALL_PAIR_JAST_HEADER

#include<string>
#include<vector>
#include<complex>

#include<formic/exception.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/wfn/configuration.h>

namespace formic {

  class InputBase;
  class Archive;
  template<class S> class Wavefunction;

  void init_all_pair_jast_commands();

  namespace fqmc {

    ///////////////////////////////////////////////////////////////////////////////////////////////////
    /// \brief  A class for all-pair jastrow factor coefficient objects
    ///
    ///////////////////////////////////////////////////////////////////////////////////////////////////
    template<class S> class AllPairJast : public formic::CoeffObj<S> {

      private:

        int _norb;                      ///<  number of orbitals
        int _nap;                       ///<  number of alpha electrons
        int _nbp;                       ///<  number of beta  electrons
        int _nau;                       ///<  number of unoccupied alpha orbitals
        int _nbu;                       ///<  number of unoccupied beta  orbitals
        int _naa_var;                   ///<  number of alpha-alpha or beta-beta variables
        int _nab_var;                   ///<  number of alpha-beta  variables
        S * _jfaa;                      ///<  pointer to the alpha-alpha jastrow factor data
        S * _jfbb;                      ///<  pointer to the beta-beta jastrow factor data
        S * _jfab;                      ///<  pointer to the alpha-beta jastrow factor data
        std::vector<S> _djfaa;          ///<  dense array of alpha-alpha jastrow factor data
        std::vector<S> _djfbb;          ///<  dense array of  beta-beta  jastrow factor data
        std::vector<S> _djfab;          ///<  dense array of alpha-beta  jastrow factor data
        std::vector<S> _djfba;          ///<  dense array of  beta-alpha jastrow factor data
        bool _simple;                   ///<  whether to use the simple jastrow parameterization
        bool _restricted;               ///<  whether to use alpha/beta symmetric jastrows
        bool _aa_fixed;                 ///<  whether the aa jastrow factors are fixed or variable
        bool _bb_fixed;                 ///<  whether the bb jastrow factors are fixed or variable
        bool _ab_fixed;                 ///<  whether the ab jastrow factors are fixed or variable
        bool _initialized;              ///<  whether the object is initialized
        formic::Configuration _config;  ///<  the configuration for which the object is initialized 

        std::vector<S> _jfaa_data;    ///<  aa jastrow factor data (not always used)
        std::vector<S> _jfbb_data;    ///<  bb jastrow factor data (not always used)
        std::vector<S> _jfab_data;    ///<  ab jastrow factor data (not always used)

        // intermediate arrays

        std::vector<S> _Jsi_a;        ///<  occupied alpha single index intermediate
        std::vector<S> _Jsi_b;        ///<  occupied beta  single index intermediate
        std::vector<S> _Jsa_a;        ///<  unoccupied alpha single index intermediate
        std::vector<S> _Jsa_b;        ///<  unoccupied beta  single index intermediate
        std::vector<S> _Jdij_aa;      ///<  occ occ alpha alpha double index intermediate
        std::vector<S> _Jdij_bb;      ///<  occ occ beta  beta  double index intermediate
        std::vector<S> _Jdij_ab;      ///<  occ occ alpha beta  double index intermediate
        std::vector<S> _Jdab_aa;      ///<  uoc uoc alpha alpha double index intermediate
        std::vector<S> _Jdab_bb;      ///<  uoc uoc beta  beta  double index intermediate
        std::vector<S> _Jdab_ab;      ///<  uoc uoc alpha beta  double index intermediate
        std::vector<S> _Jdia_aa;      ///<  occ uoc alpha alpha double index intermediate
        std::vector<S> _Jdia_bb;      ///<  occ uoc beta  beta  double index intermediate
        std::vector<S> _Jdia_ab;      ///<  occ uoc alpha beta  double index intermediate
        std::vector<S> _Jdai_aa;      ///<  uoc occ alpha alpha double index intermediate
        std::vector<S> _Jdai_bb;      ///<  uoc occ beta  beta  double index intermediate
        std::vector<S> _Jdai_ab;      ///<  uoc occ alpha beta  double index intermediate

        std::vector<S> _dzfaa;          ///<  dense array of alpha-alpha derivative jastrow factor data
        std::vector<S> _dzfbb;          ///<  dense array of  beta-beta  derivative jastrow factor data
        std::vector<S> _dzfab;          ///<  dense array of alpha-beta  derivative jastrow factor data
        std::vector<S> _dzfba;          ///<  dense array of  beta-alpha derivative jastrow factor data

        S _Zne;                       ///<  non-excitation contraction
        std::vector<S> _Zsi_a;        ///<  occupied alpha single index intermediate
        std::vector<S> _Zsi_b;        ///<  occupied beta  single index intermediate
        std::vector<S> _Zsa_a;        ///<  unoccupied alpha single index intermediate
        std::vector<S> _Zsa_b;        ///<  unoccupied beta  single index intermediate
        std::vector<S> _Zdij_aa;      ///<  occ occ alpha alpha double index intermediate
        std::vector<S> _Zdij_bb;      ///<  occ occ beta  beta  double index intermediate
        std::vector<S> _Zdij_ab;      ///<  occ occ alpha beta  double index intermediate
        std::vector<S> _Zdab_aa;      ///<  uoc uoc alpha alpha double index intermediate
        std::vector<S> _Zdab_bb;      ///<  uoc uoc beta  beta  double index intermediate
        std::vector<S> _Zdab_ab;      ///<  uoc uoc alpha beta  double index intermediate
        std::vector<S> _Zdia_aa;      ///<  occ uoc alpha alpha double index intermediate
        std::vector<S> _Zdia_bb;      ///<  occ uoc beta  beta  double index intermediate
        std::vector<S> _Zdia_ab;      ///<  occ uoc alpha beta  double index intermediate
        std::vector<S> _Zdai_aa;      ///<  uoc occ alpha alpha double index intermediate
        std::vector<S> _Zdai_bb;      ///<  uoc occ beta  beta  double index intermediate
        std::vector<S> _Zdai_ab;      ///<  uoc occ alpha beta  double index intermediate

        AllPairJast(const AllPairJast<S> & other);

        /// \brief disable the assignment operator
        AllPairJast<S> & operator=(const AllPairJast<S> & other);

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief checks that the object has been initialized for a configuration
        ///
        /// \param[in]      calling_func  the name of the function in which this check is performed
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        void check_initialized(const std::string & calling_func) {

          if ( !_initialized )
            throw formic::Exception("%s requires the object to first be initialized") % calling_func;

        }

        void fill_dense_aa(const S * const jaa,
                                 S * const daa);

        void fill_dense_ab(const S * const jab,
                                 S * const dab,
                                 S * const dba);

        void init_single_intermed(S * const val_array,
                                  const int n,
                                  const int plus_offset,
                                  const int minus_offset,
                                  const bool do_exp,
                                  const S * const jaa,
                                  const S * const jba,
                                  const int * const o2s,
                                  const int * const begin_a,
                                  const int * const begin_b);

        void init_double_intermed(S * const val_array,
                                  const int ni,
                                  const int nj,
                                  const int ex_off_i,
                                  const int rg_off_i,
                                  const int ex_off_j,
                                  const int rg_off_j,
                                  const bool do_exp,
                                  const S * const jxx,
                                  const int * const ou2si,
                                  const int * const ou2sj);

        void init_single_intermed_simple(S * const val_array,
                                         const int n,
                                         const bool do_exp,
                                         const S * const ja,
                                         const S * const jb,
                                         const S * const js,
                                         const bool do_same_site,
                                         const int * const o2s,
                                         const int * const ao2s,
                                         const int * const bo2s,
                                         const double intmd_sign);

        void init_double_intermed_simple(S * const val_array,
                                         const int n1,
                                         const int n2,
                                         const bool do_exp,
                                         const S * const jd,
                                         const int * const o2s1,
                                         const int * const o2s2,
                                         const double intmd_sign);

        void init_intermed_detail(const S * const jaa,
                                  const S * const jbb,
                                  const S * const jab,
                                  const S * const jba,
                                        S * const Jsi_a,
                                        S * const Jsi_b,
                                        S * const Jsa_a,
                                        S * const Jsa_b,
                                        S * const Jdij_aa,
                                        S * const Jdij_bb,
                                        S * const Jdij_ab,
                                        S * const Jdab_aa,
                                        S * const Jdab_bb,
                                        S * const Jdab_ab,
                                        S * const Jdia_aa,
                                        S * const Jdia_bb,
                                        S * const Jdia_ab,
                                        S * const Jdai_aa,
                                        S * const Jdai_bb,
                                        S * const Jdai_ab,
                                  const bool do_exp);

        void init_intermediates();

        static void jastrow_transform_oei(int ni,
                                          int na,
                                          double *   Jsi,
                                          double *   Jsa,
                                          double *  Jdia,
                                          double *   Tia);

        static void jastrow_transform_oei(int ni,
                                          int na,
                                          std::complex<double> *   Jsi,
                                          std::complex<double> *   Jsa,
                                          std::complex<double> *  Jdia,
                                          std::complex<double> *   Tia);

        static void jastrow_transform_tei(int ni,
                                          int na,
                                          int nj,
                                          int nb,
                                          double *   Jsi,
                                          double *   Jsa,
                                          double *   Jsj,
                                          double *   Jsb,
                                          double *  Jdia,
                                          double *  Jdij,
                                          double *  Jdib,
                                          double *  Jdaj,
                                          double *  Jdab,
                                          double *  Jdjb,
                                          double * Tiajb);

        static void jastrow_transform_tei(int ni,
                                          int na,
                                          int nj,
                                          int nb,
                                          std::complex<double> *   Jsi,
                                          std::complex<double> *   Jsa,
                                          std::complex<double> *   Jsj,
                                          std::complex<double> *   Jsb,
                                          std::complex<double> *  Jdia,
                                          std::complex<double> *  Jdij,
                                          std::complex<double> *  Jdib,
                                          std::complex<double> *  Jdaj,
                                          std::complex<double> *  Jdab,
                                          std::complex<double> *  Jdjb,
                                          std::complex<double> * Tiajb);

        static void jastrow_der_trans_oei(int ni,
                                          int na,
                                          double     Jne,
                                          double *   Jsi,
                                          double *   Jsa,
                                          double *  Jdia,
                                          double *   Tia);

        static void jastrow_der_trans_oei(int ni,
                                          int na,
                                          std::complex<double>     Jne,
                                          std::complex<double> *   Jsi,
                                          std::complex<double> *   Jsa,
                                          std::complex<double> *  Jdia,
                                          std::complex<double> *   Tia);

        static void jastrow_der_trans_tei(int ni,
                                          int na,
                                          int nj,
                                          int nb,
                                          double     Jne,
                                          double *   Jsi,
                                          double *   Jsa,
                                          double *   Jsj,
                                          double *   Jsb,
                                          double *  Jdia,
                                          double *  Jdij,
                                          double *  Jdib,
                                          double *  Jdaj,
                                          double *  Jdab,
                                          double *  Jdjb,
                                          double * Tiajb);

        static void jastrow_der_trans_tei(int ni,
                                          int na,
                                          int nj,
                                          int nb,
                                          std::complex<double>     Jne,
                                          std::complex<double> *   Jsi,
                                          std::complex<double> *   Jsa,
                                          std::complex<double> *   Jsj,
                                          std::complex<double> *   Jsb,
                                          std::complex<double> *  Jdia,
                                          std::complex<double> *  Jdij,
                                          std::complex<double> *  Jdib,
                                          std::complex<double> *  Jdaj,
                                          std::complex<double> *  Jdab,
                                          std::complex<double> *  Jdjb,
                                          std::complex<double> * Tiajb);

        S get_exponent_sum(const formic::Configuration & config,
                           const S * const jaa,
                           const S * const jbb,
                           const S * const jab,
                           const bool include_fixed);

      protected:

        void update_variables_detail(const S * const dfc, const S alpha);

      public:

        AllPairJast(const formic::InputBase & userinp, formic::Wavefunction<S> & wfn, formic::Archive & arch);

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief returns the name of the child class
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        std::string class_name() const { return "AllPairJast"; }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief returns the number of derivative ratios, which is one plus the number of variables
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        int n_der() const {
          return   ( _aa_fixed ? 0 : _naa_var )
                 + ( ( _bb_fixed || _restricted ) ? 0 : _naa_var )
                 + ( _ab_fixed ? 0 : _nab_var );
        }

        ///////////////////////////////////////////////////////////////////////////////////////////////////
        /// \brief returns the number of variables (counting the "fixed" ones)
        ///
        ///////////////////////////////////////////////////////////////////////////////////////////////////
        int n_var() const {
          return ( _restricted ? 1 : 2 ) * _naa_var + _nab_var;
        }

        bool aa_fixed() const { return _aa_fixed; }
        bool bb_fixed() const { return _bb_fixed; }
        bool ab_fixed() const { return _ab_fixed; }

        int naa_var() const { return _naa_var; }
        int nab_var() const { return _nab_var; }

        S Zne() const { return _Zne; }

        S der_ic(const S * const dfc);

        void initialize(const formic::Configuration & config);

        void update(const formic::ConfigDiff & diff);

        formic::MantExp<S> get_coeff(const formic::Configuration & config);

        S get_ratio(const formic::ConfigDiff & diff);

        void der_ratios(S * const ratios, const bool ign_fixed = true);

        void to_vector_detail(S * const vars, const bool all);
        void to_vector(S * const vars);
        void all_to_vector(S * const vars);

        void from_vector_detail(const S * const vars, const bool all);
        void from_vector(const S * const vars);
        void all_from_vector(const S * const vars);

        void print_variables();

        void transform_ints_aa_aaaa(S * const Tia, S * const Tiajb);
        void transform_ints_bb_bbbb(S * const Tia, S * const Tiajb);
        void transform_ints_aabb(S * const Tiajb);

        void init_der_intermediates(const S * const dfc);

        void der_trans_ints_aa_aaaa(S * const Tia, S * const Tiajb);
        void der_trans_ints_bb_bbbb(S * const Tia, S * const Tiajb);
        void der_trans_ints_aabb(S * const Tiajb);

        boost::shared_ptr<formic::CoeffObj<S> > clone() const;

        void get_aa_bb_ab_position_pointers(S * const vec,
                                            S * & aa_pos,
                                            S * & bb_pos,
                                            S * & ab_pos,
                                            const bool ign_fixed = true) const;

        void get_aa_bb_ab_position_pointers(const S * const vec,
                                            const S * & aa_pos,
                                            const S * & bb_pos,
                                            const S * & ab_pos,
                                            const bool ign_fixed = true) const;

        S * Jsi_a_ptr() { return &_Jsi_a.at(0); }
        S * Jsi_b_ptr() { return &_Jsi_b.at(0); }
        S * Jsa_a_ptr() { return &_Jsa_a.at(0); }
        S * Jsa_b_ptr() { return &_Jsa_b.at(0); }
        S * Jdij_aa_ptr() { return &_Jdij_aa.at(0); }
        S * Jdij_bb_ptr() { return &_Jdij_bb.at(0); }
        S * Jdij_ab_ptr() { return &_Jdij_ab.at(0); }
        S * Jdab_aa_ptr() { return &_Jdab_aa.at(0); }
        S * Jdab_bb_ptr() { return &_Jdab_bb.at(0); }
        S * Jdab_ab_ptr() { return &_Jdab_ab.at(0); }
        S * Jdia_aa_ptr() { return &_Jdia_aa.at(0); }
        S * Jdia_bb_ptr() { return &_Jdia_bb.at(0); }
        S * Jdia_ab_ptr() { return &_Jdia_ab.at(0); }
        S * Jdai_aa_ptr() { return &_Jdai_aa.at(0); }
        S * Jdai_bb_ptr() { return &_Jdai_bb.at(0); }
        S * Jdai_ab_ptr() { return &_Jdai_ab.at(0); }

        void intermed_adj_to_der_vec_aa(const int * const ao2s,
                                        const int * const bo2s,
                                        const int * const au2s,
                                        const int * const bu2s,
                                        const S * const Jsi_a,
                                        const S * const Jsa_a,
                                        const S * const Jdij_aa,
                                        const S * const Jdab_aa,
                                        const S * const Jdia_aa,
                                        const S * const Jdai_aa,
                                        const S * const Asi_a,
                                        const S * const Asa_a,
                                        const S * const Adij_aa,
                                        const S * const Adab_aa,
                                        const S * const Adia_aa,
                                        const S * const Adai_aa,
                                              S * const der_vec_aa,
                                              S * const der_vec_ab);

        void intermed_adj_to_der_vec_ab(const int * const ao2s,
                                        const int * const bo2s,
                                        const int * const au2s,
                                        const int * const bu2s,
                                        const S * const Jsi_a,
                                        const S * const Jsa_a,
                                        const S * const Jsj_b,
                                        const S * const Jsb_b,
                                        const S * const Jdia_aa,
                                        const S * const Jdij_ab,
                                        const S * const Jdib_ab,
                                        const S * const Jdaj_ab,
                                        const S * const Jdab_ab,
                                        const S * const Jdjb_bb,
                                        const S * const Asi_a,
                                        const S * const Asa_a,
                                        const S * const Asj_b,
                                        const S * const Asb_b,
                                        const S * const Adia_aa,
                                        const S * const Adij_ab,
                                        const S * const Adib_ab,
                                        const S * const Adaj_ab,
                                        const S * const Adab_ab,
                                        const S * const Adjb_bb,
                                              S * const der_vec_aa_a,
                                              S * const der_vec_ab_a,
                                              S * const der_vec_aa_b,
                                              S * const der_vec_ab_b);

        void expand_adj_single_intermed_simple(const S * const adj_array,
                                               const S * const val_array,
                                               const int n,
                                                     S * const jaa,
                                                     S * const jab,
                                                     //S * const ja,
                                                     //S * const jb,
                                                     //S * const js,
                                               const bool do_same_site,
                                               const int * const o2s,
                                               const int * const ao2s,
                                               const int * const bo2s,
                                               const double intmd_sign);

        void expand_adj_double_intermed_simple_aa(const S * const adj_array,
                                                  const S * const val_array,
                                                  const int n1,
                                                  const int n2,
                                                        S * const jd,
                                                  const int * const o2s1,
                                                  const int * const o2s2,
                                                  const double intmd_sign);

        void expand_adj_double_intermed_simple_ab(const S * const adj_array,
                                                  const S * const val_array,
                                                  const int n1,
                                                  const int n2,
                                                        S * const jd,
                                                  const int * const o2s1,
                                                  const int * const o2s2,
                                                  const double intmd_sign);

        void dense_aa_adj_to_der_vec(S * const jaa, const S * const daa);

        void dense_ab_adj_to_der_vec(S * const jab, const S * const dab, const S * const dba);

        void add_beta_derivs_to_alpha_derivs(S * const afirst, const S * const bfirst);

    };

  }

}

#endif
